import time
import numpy as np
import collections
import matplotlib.pyplot as plt
import dm_env
from pyquaternion import Quaternion

import rospy
from aloha_scripts.arm_control.msg import JointControl
from aloha_scripts.arm_control.msg import JointInformation
# from constants import DT, START_ARM_POSE
from aloha_scripts.robot_utils import Recorder, ImageRecorder
# Left finger position limits (qpos[7]), right_finger = -1 * left_finger
MASTER_GRIPPER_POSITION_OPEN = 1
MASTER_GRIPPER_POSITION_CLOSE = 0.0
PUPPET_GRIPPER_POSITION_OPEN = 1
PUPPET_GRIPPER_POSITION_CLOSE = 0.0

# Gripper joint limits (qpos[6])
MASTER_GRIPPER_JOINT_OPEN = 1
MASTER_GRIPPER_JOINT_CLOSE = 0
PUPPET_GRIPPER_JOINT_OPEN = 1
PUPPET_GRIPPER_JOINT_CLOSE = 0
############################ Helper functions ############################

MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE
PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE
MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x))

MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))

MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)

MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))

MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2

class PVRecorder:
    def __init__(self, side, init_node=True, is_debug=False):
        from collections import deque
        import rospy

        self.secs = None
        self.nsecs = None
        self.qpos = None
        self.effort = None
        self.arm_command = None
        self.gripper_command = None
        self.is_debug = is_debug

        if side == 'left':
            rospy.Subscriber("/joint_information", JointInformation, self.callback_ji)
        if side == 'right':
            rospy.Subscriber("/joint_information2", JointInformation, self.callback_ji)

        if self.is_debug:
            self.joint_timestamps = deque(maxlen=50)
            self.arm_command_timestamps = deque(maxlen=50)
            self.gripper_command_timestamps = deque(maxlen=50)
        time.sleep(0.1)

    def callback_ji(self,data):
        self.qpos = data.joint_pos
        self.qvel = data.joint_vel
        self.effort = data.joint_cur
        self.data = data
        if self.is_debug:
            self.joint_timestamps.append(time.time())

class PVEnv:
    def __init__(self, init_node, setup_robots=True, setup_base=False):
        rospy.init_node('PVEnv', anonymous=True)
        self.recorder_left = PVRecorder('left', init_node=False)
        self.recorder_right = PVRecorder('right', init_node=False)
        self.image_recorder = ImageRecorder(init_node=False)
        self.action = np.zeros(14)
        rospy.Subscriber("/joint_control",  JointControl, self.callback_jc1)
        rospy.Subscriber("/joint_control2", JointControl, self.callback_jc2)
        self.pub = rospy.Publisher('/pv_joint_control', JointControl, queue_size=10)
        self.pub2 = rospy.Publisher('/pv_joint_control2', JointControl, queue_size=10)
        # self.gripper_command = JointSingleCommand(name="gripper")

    def reset(self, fake=False):
        # if not fake:
            # # Reboot puppet robot gripper motors
            # self.puppet_bot_left.dxl.robot_reboot_motors("single", "gripper", True)
            # self.puppet_bot_right.dxl.robot_reboot_motors("single", "gripper", True)
            # self._reset_joints()
            # self._reset_gripper()
        return dm_env.TimeStep(
            step_type=dm_env.StepType.FIRST,
            reward=self.get_reward(),
            discount=None,
            observation=self.get_observation())
    def get_qpos(self):
        left_qpos_raw = self.recorder_left.qpos
        right_qpos_raw = self.recorder_right.qpos
        left_arm_qpos = left_qpos_raw[:6]
        right_arm_qpos = right_qpos_raw[:6]
        left_gripper_qpos = [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(left_qpos_raw[6])]  # this is position not joint
        right_gripper_qpos = [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(right_qpos_raw[6])]  # this is position not joint
        return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos])


    def get_qvel(self):
        left_qvel_raw = self.recorder_left.qvel
        right_qvel_raw = self.recorder_right.qvel
        left_arm_qvel = left_qvel_raw[:6]
        right_arm_qvel = right_qvel_raw[:6]
        left_gripper_qvel = [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(left_qvel_raw[6])]
        right_gripper_qvel = [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(right_qvel_raw[6])]
        return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel])


    def get_effort(self):
        left_effort_raw = self.recorder_left.effort
        right_effort_raw = self.recorder_right.effort
        left_robot_effort = left_effort_raw[:7]
        right_robot_effort = right_effort_raw[:7]
        return np.concatenate([left_robot_effort, right_robot_effort])

    def get_base_vel(self):
        base_linear_vel = 0
        base_angular_vel = 0
        return np.array([base_linear_vel, base_angular_vel])

    def get_observation(self, get_tracer_vel=False):
        obs = collections.OrderedDict()
        obs['qpos'] = self.get_qpos()
        obs['qvel'] = self.get_qvel()
        obs['effort'] = self.get_effort()
        obs['images'] = self.get_images()
        # obs['base_vel_t265'] = self.get_base_vel_t265()
        obs['base_vel'] = self.get_base_vel()
        if get_tracer_vel:
            obs['tracer_vel'] = self.get_tracer_vel()
        return obs

    def get_reward(self):
        return 0

    def get_tracer_vel(self):
        linear_vel, angular_vel = 0,0
        return np.array([linear_vel, angular_vel])
    def step(self, action, base_action=None, get_tracer_vel=False, get_obs=True):
        left_action = action[:7]
        right_action = action[7:]
        self.left_arm_control(left_action)
        self.right_arm_control(right_action)
        # if base_action is not None:
            # linear_vel_limit = 1.5
            # angular_vel_limit = 1.5
            # base_action_linear = np.clip(base_action[0], -linear_vel_limit, linear_vel_limit)
            # base_action_angular = np.clip(base_action[1], -angular_vel_limit, angular_vel_limit)
            # base_action_linear, base_action_angular = base_action
            # self.tracer.SetMotionCommand(linear_vel=base_action_linear, angular_vel=base_action_angular)
        # time.sleep(DT)
        if get_obs:
            obs = self.get_observation(get_tracer_vel)
        else:
            obs = None
        return dm_env.TimeStep(
            step_type=dm_env.StepType.MID,
            reward=self.get_reward(),
            discount=None,
            observation=obs)

    def get_action(self):
        return self.action

    def callback_jc1(self,data):
        # self.pub.publish(data)
        self.action[:6] = list(data.joint_pos[:6])
        self.action[6] = MASTER_GRIPPER_JOINT_NORMALIZE_FN(data.joint_pos[6])

    def callback_jc2(self,data):
        # self.pub2.publish(data)
        self.action[7:7+6] = list(data.joint_pos[:6])
        self.action[7+6] = MASTER_GRIPPER_JOINT_NORMALIZE_FN(data.joint_pos[6])

    def left_arm_control(self, action):
        data = JointControl()
        for i in range(7):
            data.joint_pos[i] = action[i]
        self.pub.publish(data)

    def right_arm_control(self, action):
        data = JointControl()
        for i in range(7):
            data.joint_pos[i] = action[i]
        self.pub2.publish(data)

    def get_images(self):
        return self.image_recorder.get_images()


def make_real_env(init_node, setup_robots=True, setup_base=False):
    env = PVEnv(init_node, setup_robots, setup_base)
    return env
